'''
FilePath: \src\model\ga.py
Description: 

Author: xy
Date: 2025-06-11 15:43:01
Version: 
'''
# 标准库
import random
import copy

# 第三方库

# 自定义库
from model.parallel_reverse import decode_chromosome,resolve_conflicts
from utils.plot import plot_gantt_chart

class GA:
    '''遗传算法类'''
    def __init__(self,config,info):
        '''
            num_stage,list_cast_info,self.start_index,self.list_stage,self.M_bs
        '''
        self.info=info
        ga_info=info.get_info_ga()
        self.num_stage=ga_info[0]# 阶段数量
        self.list_cast_info=ga_info[1]# 浇次信息
        self.start_index=ga_info[2]
        self.list_stage=ga_info[3]# 浇次 b 的生产阶段集合
        self.M_bs=ga_info[4]# 浇次b在阶段s可选机器集合
        self.population_size =config.getint("ga","population_size")  # 种群大小
        self.num_generations =config.getint("ga","num_generations")  # 迭代次数
        self.mutation_rate =config.getfloat("ga","mutation_rate")  # 变异概率
        self.crossover_rate =config.getfloat("ga","crossover_rate")  # 交叉概率
        self.stage_mu = {}  # 初始化空字典
        self.m = [[] for _ in range(self.num_stage)]
        for i in range(len(self.M_bs)):
            for j in range(len(self.M_bs[i])):
                if len(self.M_bs[i][j]) > 1:
                    if j not in self.stage_mu:  # 检查j是否已经在stage中
                        self.stage_mu[j] = []  # 如果不在，初始化一个空列表
                    self.stage_mu[j].append(i)
                    for k in self.M_bs[i][j]:
                        if k not in self.m[j]:  # 检查j是否已经在stage中  # 如果不在，初始化一个空列表
                            self.m[j].append(k) 


        self.stage_cr={}
        for i in range(len(self.M_bs)):
            for j in range(len(self.M_bs[i])):
                if len(self.M_bs[i][j]) > 1:
                    if i not in self.stage_cr:  # 检查j是否已经在stage中
                        self.stage_cr[i] = []  # 如果不在，初始化一个空列表
                    self.stage_cr[i].append(j)
        print(self.stage_cr)
        print(self.stage_mu)
    

    def initialize_population(self):   
        """初始化种群
        Args:
            
        Returns:
            populations (list): 种群列表
        """
        list_populations = []
        for i in range(self.population_size):
            population = []
            for j in range(len(self.list_cast_info)):
                for k in range(self.list_cast_info[j][0]):  # 浇次
                    pop = []
                    for l in range(self.num_stage):  # 阶段
                        if l in self.list_cast_info[j][1]:  # 可选阶段
                            pop.append(random.choice(self.list_cast_info[j][2][l])+1)  # 随机选择一个数
                        else:  # 不可选阶段
                            pop.append(0)
                    population.append(pop)
            list_populations.append(population)
        return list_populations


    def local_search(self,pop,stage,index):
        '''染色体局部搜索
        Args: 
            pop(list):单个染色体
            start_index(list):每个浇次的起始位置
            stage(dict):每个浇次的可选阶段
            index(int):1:搜索方式1 2:搜索方式2

        Returns:
            population(list):单个染色体
        '''
        # 染色体内部交叉
        population=copy.deepcopy(pop)
        keys=list(stage.keys())
        i=random.choice(keys)
        selected_range =copy.deepcopy(population[self.start_index[i]:self.start_index[i+1]]) 
        if len(selected_range) >= 2:
            if(index==1):
                #选择两个位置相近的染色体(距离不超过5)
                idx1 = random.randint(0, len(selected_range)-4)
                idx2 = min((idx1 + random.randint(1,4)),len(selected_range)-1)
                crossover_point = random.choice(stage[i])
                if selected_range[idx1][crossover_point]!= selected_range[idx2][crossover_point]:
                    # 交换交叉点的基因
                        selected_range[idx1][crossover_point], selected_range[idx2][crossover_point] = selected_range[idx2][crossover_point], selected_range[idx1][crossover_point]     
            if(index==2):       
                idx1 = random.randint(1, len(selected_range)-2)
                idx2 = idx1 +1
                idx3 = idx1 -1
                j = random.choice(stage[i])
                if selected_range[idx1][j]!= selected_range[idx2][j] and selected_range[idx1][j]== selected_range[idx3][j]:  # 确保两个位置的基因不同
                    crossover_point = j  # 选择交叉点
                        # 交换交叉点的基因
                    selected_range[idx1][crossover_point], selected_range[idx2][crossover_point] = selected_range[idx2][crossover_point], selected_range[idx1][crossover_point]                                        
                if selected_range[idx1][j]!= selected_range[idx3][j] and selected_range[idx1][j]== selected_range[idx2][j]:  # 确保两个位置的基因不同

                    crossover_point = j  # 选择交叉点
                        # 交换交叉点的基因
                    selected_range[idx1][crossover_point], selected_range[idx3][crossover_point] = selected_range[idx3][crossover_point], selected_range[idx1][crossover_point]



        population[self.start_index[i]:self.start_index[i+1]] = selected_range
    
        return population


    def cross_pairs(self,list_populations,list_fitness):

        new_list_populations=[]
        # 获取种群最优/最差适应度用于自适应概率
        min_fitness = min(list_fitness)
        max_fitness = max(list_fitness)
        
        while len(new_list_populations)<len(list_populations):
            idx1, idx2 = random.sample(range(len(list_populations)), 2)
            # 自适应交叉概率（基于个体适应度差异）
            fit_diff = abs(list_fitness[idx1] - list_fitness[idx2])
            current_cross_prob = self.crossover_rate * (1 - (fit_diff / (max_fitness - min_fitness + 1e-6)))  # 差异越小，交叉概率越高
            current_cross_prob = max(0.4, min(0.8, current_cross_prob))  # 限制范围
            
            if random.random() < current_cross_prob:
                # 策略1：均匀交叉（每个基因位独立交换）
                if random.random() < 0.1:  # 10%概率触发均匀交叉
                    pop1, pop2 = [], []
                    for gene1, gene2 in zip(list_populations[idx1], list_populations[idx2]):
                        # 每个基因位独立决定是否交换（均匀交叉核心）
                        pop1.append(gene1 if random.random()<0.5 else gene2)
                        pop2.append(gene2 if random.random()<0.5 else gene1)

                # if random.random() < 0.2:
                #     # 策略2：两点交叉
                #     # 生成多个交叉点
                #     crossover_points = sorted(random.sample(range(1, len(list_populations[0])), 2))
                #     pop1 = list_populations[idx1][:crossover_points[0]] + \
                #         list_populations[idx2][crossover_points[0]:crossover_points[1]] + \
                #         list_populations[idx1][crossover_points[1]:]
                #     pop2 = list_populations[idx2][:crossover_points[0]] + \
                #         list_populations[idx1][crossover_points[0]:crossover_points[1]] + \
                #         list_populations[idx2][crossover_points[1]:]
                    
                #     crossover_points2 = sorted(random.sample(range(1, len(pop1[0])), 2))  # 随机选取两个交叉点
                #     pop1 = [pop1[i][:crossover_points2[0]] + pop2[i][crossover_points2[0]:crossover_points2[1]] + pop1[i][crossover_points2[1]:] for i in range(len(pop1))]
                #     pop2 = [pop2[i][:crossover_points2[0]] + pop1[i][crossover_points2[0]:crossover_points2[1]] + pop2[i][crossover_points2[1]:] for i in range(len(pop1))]
                    
                else:
                        # 新增策略3：顺序交叉（Order Crossover, OX）
                        # 随机选择两个交叉点（确保start < end）
                    start, end = sorted(random.sample(range(len(list_populations[idx1])), 2))
                        
                    # 提取父代中间段
                    parent1 = list_populations[idx1]
                    parent2 = list_populations[idx2]
                    middle1 = parent1[start:end]  # 父代1的中间段
                    middle2 = parent2[start:end]  # 父代2的中间段
                    
                    # 构建后代1：保留父代1的中间段，按顺序填充父代2的非中间段基因
                    child1 = [None] * len(parent1)
                    child1[start:end] = middle1  # 填充中间段
                    ptr = end % len(parent1)  # 指针从end位置开始循环
                    for gene in parent2[end:] + parent2[:end]:  # 按顺序遍历父代2（从end到末尾+开头到end）
                        if child1[ptr] is None:
                            child1[ptr] = gene
                            ptr = (ptr + 1) % len(parent1)
                    
                    # 构建后代2：保留父代2的中间段，按顺序填充父代1的非中间段基因
                    child2 = [None] * len(parent2)
                    child2[start:end] = middle2  # 填充中间段
                    ptr = end % len(parent2)  # 指针从end位置开始循环
                    for gene in parent1[end:] + parent1[:end]:  # 按顺序遍历父代1（从end到末尾+开头到end）
                        if child2[ptr] is None:
                            child2[ptr] = gene
                            ptr = (ptr + 1) % len(parent2)
                        
                    pop1, pop2 = child1, child2
                    
                    new_list_populations.append(pop1)
                    new_list_populations.append(pop2)

            else:
                new_list_populations.append(list_populations[idx1])
                new_list_populations.append(list_populations[idx2])

        return new_list_populations
    
    
    def mutation(self,list_populations,stage,m,mutation_rate):

        keys = list(stage.keys())
        for time in range(int(len(list_populations))):
            population = list_populations[time]
            # 动态调整变异概率
             # 初始化 freq 字典
             
            for mutation_point in keys:  # 遍历每个变异点
                freq = {} 
                # mutation_point = random.choice(keys)
                for k in range(len(population)):  # 炉次
                    
                    if population[k][mutation_point] !=0:
                        if population[k][mutation_point] in freq:
                            freq[population[k][mutation_point]] += 1
                        else:
                            freq[population[k][mutation_point]] = 1

                for key in freq:
                    # 使用指数放大差异
                    freq[key] = freq[key] ** 2.5
                # 归一化处理
                total = sum(freq.values())
                for key in freq:
                    freq[key] = freq[key] / total

                for i in stage[mutation_point]:  # 浇次
                    for j in range(self.start_index[i], self.start_index[i + 1]):
                        if population[j][mutation_point] !=0:
                            rand_val = random.random()  # 生成一个随机数
                            if random.random() < mutation_rate:                          
                                if rand_val < freq[population[j][mutation_point]]:
                                    available_choices = [x for x in self.M_bs[i][mutation_point] if (x + 1) != population[j][mutation_point]]

                                    # available_choices = [x for x in M_bs[i][mutation_point]]
                                    if available_choices:  # 确保有可选值
                                        population[j][mutation_point] = random.choice(available_choices) + 1
        return list_populations 


    def calculate_dynamic_mutation_rate(self, current_iteration, max_iterations):
        '''动态变异率'''

        return self.mutation_rate - (self.mutation_rate - 0.01) * (current_iteration / max_iterations)


    def fitness_function(self,resolved_time_table, process_time, transport_time):
        '''
        description: 适应度函数计算
        param {*} populations: 种群
        param {*} fitness: 适应度列表
        return {*} 选择后的种群和对应的适应度数组
        '''
        waiting_time = 0
        for i in range(len(self.start_index)-1):  # 浇次
            stage=self.list_stage[i]# 阶段
            for j in range(self.start_index[i],self.start_index[i+1]):
                for s in range(len(stage)-2,-1,-1):
                    waiting_time += resolved_time_table[j][stage[s+1]]-resolved_time_table[j][stage[s]] -process_time[j][stage[s]]-transport_time[j][stage[s+1]]
    
        return waiting_time   


    def selection(self,list_populations, fitness,population_size):
        '''
        description: 选择函数(选取前1/2个体)
        param {*} populations: 种群
        param {*} fitness: 适应度列表
        return {*} 选择后的种群和对应的适应度数组
        '''
        # 将种群和适应度配对
        paired = list(zip(list_populations, fitness))
        # 按适应度排序
        paired.sort(key=lambda x: x[1])
        # 选取前1/2的个体
        selected_pairs = paired[:population_size//2]
        # 分离种群和适应度
        population1 = [p[0] for p in selected_pairs]
        fitness1 = [p[1] for p in selected_pairs]
        # 返回新种群和对应的适应度数组
        return population1, fitness1  


    def tournament_selection(self,list_populations, fitness,population_size):
        '''
        description: 锦标赛选择函数
        param {*} populations: 种群
        param {*} fitness: 适应度列表
        return {*} 选择后的种群
        '''

        sorted_indices = sorted(range(len(list_populations)), key=lambda x: fitness[x])
        selected_population = [list_populations[i] for i in sorted_indices[:10]]
        fitness1 = [fitness[i] for i in sorted_indices[:10]]
        selected_population = []
        fitness1=[]
        tournament_size = 2  # 锦标赛规模
        
        # 随机选择索引，提高效率
        candidate_indices = list(range(len(list_populations)))
        
        while len(selected_population) < population_size:
            # 随机选择 tournament_size 个个体
            tournament_indices = random.sample(candidate_indices, tournament_size)
            # 选择适应度最好的个体
            winner = min(tournament_indices, key=lambda x: fitness[x])
            selected_population.append(list_populations[winner])
            fitness1.append(fitness[winner])
        
        return selected_population,fitness1
    

    def get_load_status(self,population):
        '''获取各阶段机器均衡程度
        '''
        #获取所有机器的使用次数
        keys = list(self.stage_mu.keys())
        freqs={}
        for i in keys:
            freq={}
            for k in range(len(population)):  # 炉次
                if population[k][i] != 0:
                    if population[k][i] in freq:
                        freq[population[k][i]] += 1
                    else:
                        freq[population[k][i]] = 1
            freqs[i]=freq

        #遍历
        stage=copy.deepcopy(self.stage_cr)
        for key,value in stage.items():
            for i in value:
                machine_state=[]
                for j in self.M_bs[key][i]:
                    machine_state.append(freqs[i][j+1])
                if max(machine_state) - min(machine_state) > 5:
                    value.remove(i)
        return stage
                    

    def run(self):
        '''遗传算法主函数'''

        # 初始化种群
        list_populations = self.initialize_population()
        best_fitness = 1000000
        best_solution = None
        fitness_values = []
        for i in range(len(list_populations)):
            population = list_populations[i]
            pro_timetable, trans_timetable, start_time = self.info.get_info(population) # 加工时间，移动时间，开始时间
            time_table = decode_chromosome(population, pro_timetable, trans_timetable, start_time, self.list_stage,self.start_index)# 解码
            resolved_time_table, machine_usage = resolve_conflicts(time_table, pro_timetable, trans_timetable, population,self.info)# 冲突解决
            wait_time = self.fitness_function(resolved_time_table, pro_timetable, trans_timetable)# 适应度函数
            fitness_values.append(wait_time)        
            # 记录最优解
            if wait_time < best_fitness:
                best_fitness = wait_time
                best_solution = population
        for generation in range(self.num_generations):
            # 初始解析
            # 选择
            # if generation>50:
            #     self.population_size=300
            populations1 ,fitness1= self.selection(list_populations, fitness_values,self.population_size)

            # #锦标赛选择
            # populations1,fitness1 =self.tournament_selection(list_populations, fitness_values,self.population_size)
            populations2 = copy.deepcopy(populations1)

            # 交叉
            populations2 = self.cross_pairs(populations2, fitness1)
            # 变异
            # 假设这里有 current_iteration 和 max_iterations 变量
            dynamic_mutation_rate = self.calculate_dynamic_mutation_rate(generation, self.num_generations)
            populations2 = self.mutation(populations2, self.stage_mu, self.m,dynamic_mutation_rate)
        
            list_populations=populations1+populations2
            fitness_values = []
            
            for i in range(len(list_populations)):
                
                #读取每个染色体相应的适应度函数
                population = list_populations[i]
                pro_timetable, trans_timetable, start_time = self.info.get_info(population) # 加工时间，移动时间，开始时间
                time_table = decode_chromosome(population, pro_timetable, trans_timetable, start_time, self.list_stage,self.start_index)# 解码
                resolved_time_table, machine_usage = resolve_conflicts(time_table, pro_timetable, trans_timetable, population, self.info)# 冲突解决
                wait_time = self.fitness_function(resolved_time_table, pro_timetable, trans_timetable)# 适应度函数
                
                if generation>0:
                    #对后半段染色体采用局部搜索
                    if i>=len(list_populations)//2:
                        # stage=self.get_load_status(population)
                        # pop=self.local_search(population,stage,2)
                        #特殊操作
                        pop=self.local_search(population,self.stage_cr,2)
                        pro_timetable, trans_timetable, start_time = self.info.get_info(pop) # 加工时间，移动时间，开始时间
                        time_table = decode_chromosome(pop, pro_timetable, trans_timetable, start_time, self.list_stage, self.start_index)# 解码
                        resolved_time_table, machine_usage = resolve_conflicts(time_table, pro_timetable, trans_timetable, pop, self.info)# 冲突解决
                        wait_time1 = self.fitness_function(resolved_time_table, pro_timetable, trans_timetable)# 适应度函数
                        if wait_time1<wait_time:
                            wait_time=wait_time1
                            list_populations[i]=copy.deepcopy(pop)

                fitness_values.append(wait_time)
                
                # 记录最优解
                if wait_time < best_fitness:
                    best_fitness = wait_time
                    best_solution = population

           
            # 输出当前代信息
            print(f'Generation {generation}, Best Fitness: {best_fitness}')

        
        # 输出最终结果
        print(f'Best Solution Found: Fitness = {best_fitness}')
        pro_timetable, trans_timetable, start_time = self.info.get_info(best_solution)
        time_table = decode_chromosome(best_solution, pro_timetable, trans_timetable, start_time, self.list_stage, self.start_index)
        resolved_time_table, machine_usage = resolve_conflicts(time_table, pro_timetable, trans_timetable,best_solution, self.info)
        plot_gantt_chart(machine_usage)




